package com.sparrow.common.util;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.TreeMap;
import java.util.SortedMap;

/**
 * 一致性哈希算法实现类(带虚拟节点)，横向扩展节点时只有少部分节点会变。
 * 一致性哈希的关键在于 顺时针查找第一个大于等于目标哈希值的节点
 * 带有虚拟节点机制，用来提升数据分布的均匀性
 * 查找逻辑：使用 key 的 hash 查找其顺时针最近的节点；若超过最大值则从环头开始
 * 查找“比某个 key 大的最小 key”（tailMap() 和 firstKey()）
 */
public class ConsistentHashing {
    // 哈希环，存储虚拟节点的哈希值和对应的实际节点
    private final TreeMap<Integer, String> circle = new TreeMap<>();
    // 每个实际节点对应的虚拟节点数量
    private final int virtualNodes;

    /**
     * 构造函数，初始化一致性哈希环。
     *
     * @param virtualNodes 每个实际节点对应的虚拟节点数量
     */
    public ConsistentHashing(int virtualNodes) {
        this.virtualNodes = virtualNodes;
    }

    /**
     * 为指定的实际节点创建虚拟节点并添加到哈希环中。
     *
     * @param node 实际节点名称
     */
    public void addNode(String node) {
        for (int i = 0; i < virtualNodes; i++) {
            String virtualNodeName = createVirtualNodeName(node, i);
            int hash = hash(virtualNodeName);
            circle.put(hash, node);
        }
    }

    /**
     * 根据实际节点名称和虚拟节点编号创建虚拟节点名称。
     *
     * @param node 实际节点名称
     * @param index 虚拟节点编号
     * @return 虚拟节点名称
     */
    private String createVirtualNodeName(String node, int index) {
        return node + "&&VN" + index;
    }

    /**
     * 根据键获取对应的节点。
     *
     * @param key 键
     * @return 对应的节点名称，如果哈希环为空则返回 null
     */
    public String getNode(String key) {
        if (circle.isEmpty()) {
            return null;
        }
        int hash = hash(key);
        return findNodeByHash(hash);
    }

    /**
     * 根据哈希值查找对应的节点。大于等于 hash值的 第一个节点
     *
     * @param hash 哈希值
     * @return 对应的节点名称
     */
    private String findNodeByHash(int hash) {
        if (!circle.containsKey(hash)) {
            SortedMap<Integer, String> tailMap = circle.tailMap(hash);
            hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
        }
        return circle.get(hash);
    }

    /**
     * 计算字符串的哈希值，先计算字符串的MD5值然后再提取出hash值
     *
     * @param key 要计算哈希值的字符串
     * @return 哈希值
     */
    private int hash(String key) {
        try {
            MessageDigest md = MessageDigest.getInstance("MD5");
            byte[] digest = md.digest(key.getBytes());
            return extractHashValue(digest);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("MD5 algorithm not available", e);
        }
    }

    /**
     * 从字节数组中提取哈希值。
     *
     * @param digest 字节数组
     * @return 哈希值
     */
    private int extractHashValue(byte[] digest) {
        int h = 0;
        for (int i = 0; i < 4; i++) {
            // 左移 8 位，为下一个字节腾出位置
            h <<= 8;
            // 将当前字节和 h 进行按位或操作，合并当前字节到 h 中
            h |= ((int) digest[i]) & 0xFF;
        }
        return h;
    }

    public static void main(String[] args) {
        //每个节点映射100个虚拟节点
        ConsistentHashing ch = new ConsistentHashing(100);
        //3个实际节点
        for (int i = 1; i <= 4; i++) {
            ch.addNode("Node-" + i);
        }
        for (int i = 1; i <= 10; i++) {
            String key = "Key-" + i;
            String node = ch.getNode(key);
            System.out.println(key + " is mapped to node: " + node);
        }
    }
}    